{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## AutoGen Studio Agent Workflow API Example\n",
                "\n",
                "This notebook focuses on demonstrating capabilities of the autogen studio workflow python api.  \n",
                "\n",
                "- Declarative Specification of an Agent Team\n",
                "- Loading the specification and running the resulting agent\n",
                "\n",
                " "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from autogenstudio.teammanager import TeamManager\n",
                "\n",
                "wm = TeamManager()\n",
                "result = await wm.run(task=\"What is the weather in New York?\", team_config=\"team.json\")\n",
                "print(result)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result_stream =  wm.run_stream(task=\"What is the weather in New York?\", team_config=\"team.json\")\n",
                "async for response in result_stream:\n",
                "    print(response)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": []
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Load Directly Using the AgentChat API "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import json \n",
                "from autogen_agentchat.teams import BaseGroupChat\n",
                "team_config = json.load(open(\"team.json\"))  \n",
                "team = BaseGroupChat.load_component(team_config)\n",
                "print(team._participants)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## AutoGen Studio Database API\n",
                "\n",
                "Api for creating objects and serializing to a database."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from autogenstudio.database import DatabaseManager\n",
                "import os\n",
                "# delete database\n",
                "# if os.path.exists(\"test.db\"):\n",
                "#     os.remove(\"test.db\")\n",
                "\n",
                "os.makedirs(\"test\", exist_ok=True)\n",
                "# create a database\n",
                "dbmanager = DatabaseManager(engine_uri=\"sqlite:///test.db\", base_dir=\"test\")\n",
                "dbmanager.initialize_database()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Sample AgentChat Example (Python)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "from autogen_agentchat.agents import AssistantAgent\n",
                "from autogen_agentchat.conditions import TextMentionTermination, MaxMessageTermination\n",
                "from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat\n",
                "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
                "\n",
                "planner_agent = AssistantAgent(\n",
                "    \"planner_agent\",\n",
                "    model_client=OpenAIChatCompletionClient(model=\"gpt-4\"),\n",
                "    description=\"A helpful assistant that can plan trips.\",\n",
                "    system_message=\"You are a helpful assistant that can suggest a travel plan for a user based on their request. Respond with a single sentence\",\n",
                ")\n",
                "\n",
                "local_agent = AssistantAgent(\n",
                "    \"local_agent\",\n",
                "    model_client=OpenAIChatCompletionClient(model=\"gpt-4\"),\n",
                "    description=\"A local assistant that can suggest local activities or places to visit.\",\n",
                "    system_message=\"You are a helpful assistant that can suggest authentic and interesting local activities or places to visit for a user and can utilize any context information provided. Respond with a single sentence\",\n",
                ")\n",
                "\n",
                "language_agent = AssistantAgent(\n",
                "    \"language_agent\",\n",
                "    model_client=OpenAIChatCompletionClient(model=\"gpt-4\"),\n",
                "    description=\"A helpful assistant that can provide language tips for a given destination.\",\n",
                "    system_message=\"You are a helpful assistant that can review travel plans, providing feedback on important/critical tips about how best to address language or communication challenges for the given destination. If the plan already includes language tips, you can mention that the plan is satisfactory, with rationale.Respond with a single sentence\",\n",
                ")\n",
                "\n",
                "travel_summary_agent = AssistantAgent(\n",
                "    \"travel_summary_agent\",\n",
                "    model_client=OpenAIChatCompletionClient(model=\"gpt-4\"),\n",
                "    description=\"A helpful assistant that can summarize the travel plan.\",\n",
                "    system_message=\"You are a helpful assistant that can take in all of the suggestions and advice from the other agents and provide a detailed tfinal travel plan. You must ensure th b at the final plan is integrated and complete. YOUR FINAL RESPONSE MUST BE THE COMPLETE PLAN. When the plan is complete and all perspectives are integrated, you can respond with TERMINATE.Respond with a single sentence\",\n",
                ")\n",
                "\n",
                "termination = TextMentionTermination(\"TERMINATE\") | MaxMessageTermination(10)\n",
                "group_chat = RoundRobinGroupChat(\n",
                "    [planner_agent, local_agent, language_agent, travel_summary_agent], termination_condition=termination\n",
                ")\n",
                "\n",
                "\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "result = group_chat.run_stream(task=\"Plan a 3 day trip to Nepal.\")\n",
                "async for response in result:\n",
                "    print(response)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import json\n",
                "\n",
                "# convert to config \n",
                "config = group_chat.dump_component().model_dump()\n",
                "# save as json \n",
                "\n",
                "with open(\"travel_team.json\", \"w\") as f:\n",
                "    json.dump(config, f, indent=4)\n",
                "\n",
                "# load from json\n",
                "with open(\"travel_team.json\", \"r\") as f:\n",
                "    config = json.load(f)\n",
                "\n",
                "group_chat = RoundRobinGroupChat.load_component(config) \n",
                "result = group_chat.run_stream(task=\"Plan a 3 day trip to Nepal.\") \n",
                "async for response in result:\n",
                "    print(response)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": []
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "agnext",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.11.9"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}
